import cv2
import time

from yolo import YOLO
from PIL import Image
import numpy as np


video_save_path = ""

model = YOLO(cuda=True)
capture = cv2.VideoCapture(0)

ref, frame = capture.read()
if not ref:
    raise ValueError("未能正确读取摄像头（视频），请注意是否正确安装摄像头（是否正确填写视频路径）。")

fps = 0.0
# noinspection PyRedundantParentheses
while True:
    t1 = time.time()
    # 读取某一帧
    ref, frame = capture.read()
    if not ref:
        break
    # 格式转变，BGRtoRGB
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    # 转变成Image
    frame = Image.fromarray(np.uint8(frame))
    # 进行检测
    frame = np.array(model.detect_image(frame))
    # RGBtoBGR满足opencv显示格式
    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

    fps = (fps + (1. / (time.time() - t1))) / 2
    # noinspection PyRedundantParentheses
    # print("fps= %.2f" % (fps))
    # noinspection PyRedundantParentheses
    frame = cv2.putText(frame, "fps= %.2f" % (fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

    cv2.imshow("video", frame)
    c = cv2.waitKey(1) & 0xff

    if c == 27:
        capture.release()
        break

print("视频检测完毕！！！")
capture.release()
cv2.destroyAllWindows()